#!/usr/bin/python3
'''
Author: January
Date: 2021-11-21 19:49:09
'''

import socket
import logging
import time
import argparse
from mbedtls import tls
from mbedtls import x509
from mbedtls import pk
from contextlib import suppress
import random

# problem
# ECC指定曲线无效
# ECC加载PEM类型的密钥存在问题

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("dtls")

# addrs = {
#     0 : ("192.168.3.186", 9000),
#     1 : ("192.168.3.108", 9001),
#     2 : ("127.0.0.1", 9002)
# }

addrs = {
    0 : ("127.0.0.1", 9000),
    1 : ("127.0.0.1", 9001),
    2 : ("127.0.0.1", 9002)
}



def rand_bytes(length):
    return b''.join([int.to_bytes(random.randint(0,255), 1, 'little') for i in range(length)])

class GroudStation:
    def __init__(self, id, cafile, certfile, keyfile) -> None:
        udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        addr = addrs[id]
        trust_store = tls.TrustStore()
        ca_crt = x509.CRT.from_file(cafile)
        gs_crt = x509.CRT.from_file(certfile)
        trust_store.add(ca_crt)
        with open(keyfile, "rb") as f:
            keydata = f.read()
            key = pk.ECC.from_DER(keydata)
        
        tls_ctx = tls.ClientContext(tls.DTLSConfiguration(
            trust_store=trust_store,
            certificate_chain=([gs_crt, ca_crt], key),
            validate_certificates=True,
            ciphers=("TLS-ECDHE-ECDSA-WITH-AES-128-GCM-SHA256",)
         ))
         
        

        self.channel = tls_ctx.wrap_socket(udp_sock,server_hostname=None)
        self.channel.bind(addr)
    
    def set_key(self, peer_id, key):
        start_time = time.time()
        self.channel.connect(addrs[peer_id])
        self.channel.do_handshake()
        end_time = time.time()
        self.channel.send(key)
        logger.info("send key [%s]"%(key.hex()))
        logger.info("DTLSV1.2 Protocol finished in [%.2f]ms"%((end_time - start_time) * 1000))

def generate_key():
    test = pk.ECC()
    test.generate()
    with open("test_uav.key", "w") as f:
        key = test.export_key()
        print(len(key))
        f.write(key)
    return key

class Drone:
    def __init__(self, id, cafile, certfile, keyfile) -> None:
        udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        trust_store = tls.TrustStore()
        ca_crt = x509.CRT.from_file(cafile)
        uav_crt = x509.CRT.from_file(certfile)

        with open(keyfile, "rb") as f:
            keydata = f.read()
            key = pk.ECC.from_DER(keydata)
        trust_store.add(ca_crt)
        tls_ctx = tls.ServerContext(tls.DTLSConfiguration(
            trust_store=trust_store,
            certificate_chain=([uav_crt, ca_crt], key),
            validate_certificates=True,
            ciphers=("TLS-ECDHE-ECDSA-WITH-AES-128-GCM-SHA256",)
         ))
        addr = addrs[id]
        udp_sock.bind(addr)
        self.channel = tls_ctx.wrap_socket(udp_sock)
        
    def accept_key(self):
        print("waiting for connection")
        conn, addr = self.channel.accept()
        start_time = time.time()
        conn.setcookieparam(addr[0].encode())
        with suppress(tls.HelloVerifyRequest):
            conn.do_handshake()
        conn, addr = conn.accept()
        conn.setcookieparam(addr[0].encode())
        conn.do_handshake()
        end_time = time.time()
        data = conn.recv(4096)
        print("recv key [%s]"%(data.hex()))
        logger.info("DTLSv1.2 Protocol finished in [%.2f]ms"%((end_time - start_time) * 1000))

parser = argparse.ArgumentParser()
parser.add_argument("mode", nargs=1)
args = parser.parse_args()


mode = args.mode[0]
if mode == "gs":
    protocol = GroudStation(0,"cert/ca/ca.crt", "cert/gs/gs.crt", "cert/gs/gs.der")
    protocol.set_key(1, rand_bytes(16))
else:
    protocol = Drone(1, "cert/ca/ca.crt", "cert/uav/uav.crt", "cert/uav/uav.der")
    protocol.accept_key()


